#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   main.py
@Time    :   2024/11/12 16:22:59
@Author  :   XD_beta
@Version :   1.0.0
@Contact :   xd_beta@163.com
@License :   GNU GENERAL PUBLIC LICENSE Version 2, June 1991
@Desc    :   None
'''

import os, struct, hashlib, argparse
import tkinter as tk
from tkinter import filedialog, messagebox
from typing import List, Tuple

class OTAPackager:
    def __init__(self):
        self.PROTOCOL_VERSION = 0x0100
        self.PID = 0x0001  # 可配置
        self.FW_TABLE_OFFSET = 64
        self.files_info = []
        self.holes = [] # 文件表和文件之间以及文件和文件之间的偏移量
        
    def calculate_md5(self, data: bytes) -> bytes:
        """计算数据的MD5值"""
        hasher = hashlib.md5()
        hasher.update(data)
        return hasher.digest()

    def create_header(self, total_size: int, ota_version: Tuple[int, int, int]) -> bytes:
        """创建OTA包头"""
        # 将版本号转换为适当的格式 (例如: 1.3.0 -> 0x00010300)
        version_value = (ota_version[0] << 16) | (ota_version[1] << 8) | ota_version[2]
        
        # 打包头部信息
        header = struct.pack('<HHIIIB',
                           self.PROTOCOL_VERSION,    # 2 bytes
                           self.PID,                 # 2 bytes
                           version_value,            # 4 bytes
                           total_size,               # 4 bytes
                           self.FW_TABLE_OFFSET,     # 4 bytes
                           len(self.files_info))     # 1 byte
        print('<HHIIIB,{:X},{:X},{:X},{:X},{:X},{:X}'.format(self.PROTOCOL_VERSION, self.PID, version_value, total_size, self.FW_TABLE_OFFSET,len(self.files_info)))
        
        # 添加32字节的文件签名（暂时用0x00填充）
        header += b'\x00' * 32
        
        return header

    def create_fw_table_entry(
            self, major_type: int,
            minor_type: int, 
            version: Tuple[int, int, int],
            need_upgrade: int,
            fw_offset: int,
            fw_size: int,
            fw_signature: bytes) -> bytes:
        """創建固件表項"""
        # 將版本號打包為3字節
        version_bytes = bytes([
            version[0],  # major
            version[1],  # minor
            version[2]   # revision
        ])
        
        # 使用格式字符串打包前兩個字節
        entry = struct.pack('<BB', major_type, minor_type)
        # 添加3字節版本號(小端存儲)(交流判斷傳入參數是否已倒轉)
        entry += version_bytes[::-1] # 未倒轉
        # entry += version_bytes # 已倒轉
        # 打包剩餘字段
        entry += struct.pack('<BII32s',
                           need_upgrade,
                           fw_offset,
                           fw_size,
                           fw_signature)

        print('<BB,{:X},{:X},'.format(major_type, minor_type), version_bytes, ',<BII32s,{:X},{:X},{:X},'.format(need_upgrade,fw_offset,fw_size),fw_signature)
        return entry

    def pack_files(self,
                   file_paths: List[str],
                   FW_major_type: List[str],
                   FW_minor_type: List[str],
                   FW_version,
                   FW_need_upgrade,
                   output_path: str,
                   ota_version: Tuple[int, int, int] = (1, 0, 0)
        ):
        """打包文件"""
        if not file_paths:
            raise ValueError("No files selected")
        self.holes.clear()

        # 计算文件表结束位置后的偏移量
        current_offset = self.FW_TABLE_OFFSET + (len(file_paths) * 46)  # 46是每个表项的大小
        self.holes.append(0 if current_offset % 16 == 0 else 16 - current_offset % 16)
        current_offset += 0 if current_offset % 16 == 0 else 16 - current_offset % 16 # 每个文件起始位置后的偏移量 文件表和文件之间需要填充hole以对齐 使得每个文件都在行头开始 一行16bytes

        # 收集所有文件数据和信息
        all_file_data = []
        self.files_info = []
        
        for index in range(len(file_paths)):
            with open(file_paths[index], 'rb') as f:
                file_data = f.read()
            
            _ = int(FW_version[index],16)
            version = (_>>16&0xff, _>>8&0xff, _&0xff)

            file_size = len(file_data)
            file_signature = self.calculate_md5(file_data)
            
            # 存储文件信息
            self.files_info.append({
                'major_type': int(FW_major_type[index],16),
                'minor_type': int(FW_minor_type[index],16),
                'version': version,
                'need_upgrade': 1 if len(FW_need_upgrade)==0 else int(FW_need_upgrade[index],16),
                'offset': current_offset,
                'size': file_size,
                'signature': file_signature,
                'data': file_data
            })
            current_offset += file_size
            self.holes.append(0 if current_offset % 16 == 0 else 16 - current_offset % 16)
            current_offset += 0 if current_offset % 16 == 0 else 16 - current_offset % 16 # 每个文件起始位置后的偏移量 文件表和文件之间需要填充hole以对齐 使得每个文件都在行头开始 一行16bytes


        # 创建输出文件
        with open(output_path, 'wb') as out_file:
            # 写入头部
            header = self.create_header(current_offset-self.holes[-1], ota_version)
            out_file.write(header)
            
            # 写入填充
            padding_size = self.FW_TABLE_OFFSET - len(header)
            out_file.write(b'\x00' * padding_size)
            
            # 写入固件表
            for file_info in self.files_info:
                table_entry = self.create_fw_table_entry(
                    file_info['major_type'],
                    file_info['minor_type'],
                    file_info['version'],
                    file_info['need_upgrade'],
                    file_info['offset'],
                    file_info['size'],
                    file_info['signature']
                )
                out_file.write(table_entry)
            
            # 写入文件数据
            for index, file_info in enumerate(self.files_info):
                out_file.write(b'\x00' * self.holes[index])
                out_file.write(file_info['data'])
            print(self.holes)

class OTAPackagerGUI:
    def __init__(self):
        self.window = tk.Tk()
        self.window.title("OTA Packager")
        self.window.geometry("500x400")
        self.packager = OTAPackager()
        self.selected_files = []
        self.create_widgets()

    def create_widgets(self):
        # PID 输入
        tk.Label(self.window, text="PID (hex):").pack()
        self.pid_entry = tk.Entry(self.window)
        self.pid_entry.insert(0, "0001")
        self.pid_entry.pack()

        # OTA Version 输入
        tk.Label(self.window, text="OTA Version (x.x.x):").pack()
        self.version_entry = tk.Entry(self.window)
        self.version_entry.insert(0, "1.3.0")
        self.version_entry.pack()

        # major type 输入
        tk.Label(self.window, text="major type [hex]:").pack()
        self.major_type = tk.Entry(self.window)
        self.major_type.insert(0, "10,fa,dd,...")
        self.major_type.pack()

        # minor version 输入
        tk.Label(self.window, text="minor type [hex]:").pack()
        self.minor_type = tk.Entry(self.window)
        self.minor_type.insert(0, "20,af,cc,...")
        self.minor_type.pack()

        # firmware version 输入
        tk.Label(self.window, text="firmware version [hex]:").pack()
        self.firmware_version = tk.Entry(self.window)
        self.firmware_version.insert(0, "001122,221100,ff0001,...")
        self.firmware_version.pack()

        # 文件列表
        self.files_listbox = tk.Listbox(self.window, width=50, height=10)
        self.files_listbox.pack(pady=10)

        # 按钮
        tk.Button(self.window, text="Add Files", command=self.add_files).pack()
        tk.Button(self.window, text="Clear Files", command=self.clear_files).pack()
        tk.Button(self.window, text="Generate OTA Package", command=self.generate_package).pack()

    def add_files(self):
        files = filedialog.askopenfilenames(filetypes=[("Binary files", "*.bin")])
        for file in files:
            if file not in self.selected_files:
                self.selected_files.append(file)
                self.files_listbox.insert(tk.END, os.path.basename(file))

    def clear_files(self):
        self.selected_files.clear()
        self.files_listbox.delete(0, tk.END)

    def generate_package(self):
        if not self.selected_files:
            messagebox.showerror("Error", "Please select at least one file")
            return

        # 解析 PID
        pid = int(self.pid_entry.get(), 16)
        self.packager.PID = pid

        # 解析版本号
        version = tuple(map(int, self.version_entry.get().split('.')))
        if len(version) != 3:
            raise ValueError("Invalid version format")

        # 选择输出文件
        output_path = filedialog.asksaveasfilename(
            defaultextension=".bin",
            filetypes=[("Binary files", "*.bin")]
        )

        # 解析major_type minor_type firmware_version
        major_type = self.major_type.get().split(',')
        minor_type = self.minor_type.get().split(',')
        firmware_version = self.firmware_version.get().split(',')
        
        if output_path:
            self.packager.pack_files(self.selected_files, major_type, minor_type, firmware_version, [], output_path, version)
            messagebox.showinfo("Success", "OTA package generated successfully!")
        

    def run(self):
        self.window.mainloop()

class OTAPackagerCommand():
    def __init__(self):
        self.pid_hex = "0001"
        self.version = "1.0.0"
        self.files = []
        self.FW_major_type = []
        self.FW_minor_type = []
        self.FW_version = []
        self.FW_need_upgrade = []
        self.output = "./output.bin"
        self.packager = OTAPackager()
    
    def get_command(self):
        # 創建解析對象
        parser = argparse.ArgumentParser()

        # 添加該解析對象中的參數和對應的命令表示
        parser.add_argument("-P", "--pid_hex", default="0001", help="")
        parser.add_argument("-V", "--version", default="1.0.0", help="")
        parser.add_argument("-F", "--files", default=[], nargs="+", help="")
        parser.add_argument("--FW_major_type", default=[], nargs="+", help="")
        parser.add_argument("--FW_minor_type", default=[], nargs="+", help="")
        parser.add_argument("--FW_version", default=[], nargs="+", help="")
        parser.add_argument("--FW_need_upgrade", default=[], nargs="+", help="")
        parser.add_argument("-O", "--output", default="./output.bin", help="")

        self.pid_hex = parser.parse_args().pid_hex
        self.version = parser.parse_args().version
        self.files = parser.parse_args().files
        self.FW_major_type = parser.parse_args().FW_major_type
        self.FW_minor_type = parser.parse_args().FW_minor_type
        self.FW_version = parser.parse_args().FW_version
        self.FW_need_upgrade = parser.parse_args().FW_need_upgrade
        self.output = parser.parse_args().output

    # 接收整條配置
    def configure(self, configurations_hex:List[str])->List[dict]:
        configurations = []
        for item in configurations_hex:
            step = int(item, 16)
            configurations.append({
                'major_type': step >> 40 & 0xff,
                'minor_type': step >> 32 & 0xff,
                'version': step >> 8 & 0xffffff,
                'need_upgrade': step & 0xff,
            })
        return configurations

    def package(self):
        if self.files == []:
            print("Error: Please select at least one file.")
            return
        
        try:
            # 解析 PID
            pid_hex = int(self.pid_hex, 16)
            self.packager.PID = pid_hex

            # 解析版本号
            version = tuple(map(int, self.version.split('.')))
            if len(version) != 3:
                raise ValueError("Invalid version format")
            
            if self.output:
                self.packager.pack_files(
                    self.files,
                    self.FW_major_type,
                    self.FW_minor_type,
                    self.FW_version,
                    self.FW_need_upgrade,
                    self.output,
                    version)
                print("Success:OTA package generated successfully!")
        
        except Exception as e:
            print("Error:", str(e))

def by_command():
    command = OTAPackagerCommand()
    command.get_command()
    command.package()

def by_gui():
    gui = OTAPackagerGUI()
    gui.run()

''' test command
python3 main.py 
    --pid_hex fedc 
    --version 512.2.3 
    --files "./files/test_1.bin" "./files/test_2.bin" "./files/test_3.bin" 
    --FW_major_type 10 fa dd 
    --FW_minor_type 20 af cc 
    --FW_version 001122 221100 000001  
    --output "./files/output.bin"
python3 main.py 
    --pid_hex fedc 
    --version 512.2.3 
    --files "./files/test_1.bin" "./files/test_2.bin" "./files/test_3.bin" 
    --FW_major_type 10 fa dd 
    --FW_minor_type 20 af cc 
    --FW_version 001122 221100 000001 
    --FW_need_upgrade 00 00 00
    --output "./files/output.bin"
'''
if __name__ == '__main__':
    # by_command()
    by_gui()